package fos

/** Two-phase inferencer, first collect constraints, then solve them. */
class TwoPhaseInferencer extends TypeInferencers {
  import Type._

  type Constraint = (Type, Type)

  val noContraints: List[Constraint] = Nil
  case class TypingResult(tpe: Type, c: List[Constraint])

  /**
   * Type <code>t</code> in <code>env</code> and return its type and a
   *  constraint list.
   */
  def collect(env: Env, term: Term): TypingResult = term match {
    case Var(x) => {
      val t1 = lookup(env, x)
      if (t1 == null) throw TypeError("Unknown variable " + x)
      TypingResult(t1.instantiate, noContraints)
    }
    case True | False => TypingResult(TypeBool, noContraints)
    case Zero => TypingResult(TypeNat, noContraints)
    case Pred(t) => {
      val TypingResult(tp, c) = collect(env, t)
      TypingResult(TypeNat, (tp, TypeNat) :: c)
    }
    case Succ(t) => {
      val TypingResult(tp, c) = collect(env, t)
      TypingResult(TypeNat, (tp, TypeNat) :: c)
    }
    case IsZero(t) => {
      val TypingResult(tp, c) = collect(env, t)
      TypingResult(TypeBool, (tp, TypeNat) :: c)
    }
    case If(t1, t2, t3) => {
      val TypingResult(tp1, c1) = collect(env, t1)
      val TypingResult(tp2, c2) = collect(env, t2)
      val TypingResult(tp3, c3) = collect(env, t3)
      TypingResult(tp2, (tp1, TypeBool) :: (tp2, tp3) :: c1 ::: c2 ::: c3)
    }
    case Abs(v, tpTerm, t) => {
      tpTerm match {
        case EmptyType => {
          val X = Names.getFresh()
          val typesch = TypeScheme(Nil, X)
          val TypingResult(tp2, c) = collect((v, typesch) :: (X.name, typesch) :: env, t)
          TypingResult(TypeFun(X, tp2), c)
        }
        case _ => {
          val TypingResult(tp2, c) = collect((v, TypeScheme(Nil, toType(tpTerm))) :: env, t)
          TypingResult(TypeFun(toType(tpTerm), tp2), c)
        }
      }
    }
    case App(t1, t2) => {
      val TypingResult(tp1, c1) = collect(env, t1)
      val TypingResult(tp2, c2) = collect(env, t2)
      val X = Names.getFresh()
      TypingResult(X, (tp1, TypeFun(tp2, X)) :: c1 ::: c2)
    }
    case Let(x, ty, v, t) => {
      //step 1
      val TypingResult(s, c) = collect(env, v)
      //step 2
      val subs = ty match {
        case EmptyType => unify(c)
        case _ => unify(c).compose(unify((toType(ty), s) :: Nil))
      }
      val T = subs(s)
      val env2 = subs(env)
      //step 3
      val typesch = TypeScheme(listVariablesNotInEnv(env2, T), T)
      //step 4

      val env3 = (x, typesch) :: env2
      /*val env3 = ty match {
        case EmptyType => (x,typesch) :: env2
        case _ => (x,typesch) :: (x, TypeScheme(Nil, toType(ty))) :: env2
      }*/
      collect(env3, t)
      //step 5 (nothing)
    }
  }

  def listVariablesNotInEnv(env: Env, x: Type): List[TypeVar] = x match {
    case TypeVar(name) if lookup(env, name) == null => TypeVar(name) :: Nil
    case TypeFun(a, b) => listVariablesNotInEnv(env, a) ::: listVariablesNotInEnv(env, b)
    case _ => Nil
  }

  /**
   */
  def unify(c: List[Constraint]): Substitution = {
    unify(c, EmptySubst)
  }

  def unify(c: List[Constraint], s: Substitution): Substitution = {
    //   println("Substitution:"+s)
    //   println("Constraints:"+c)
    def notIn(tp1: TypeVar, tp2: Type): Boolean = tp2 match {
      case TypeVar(_) if tp1 == tp2 => false
      case TypeFun(a, b) => notIn(tp1, a) && notIn(tp1, b)
      case _ => true
    }

    def varSub(x: TypeVar, y: Type, tail: List[Constraint]) = s.findConflict(x, y) match {
      case None => s.lookup(x) match {
        case None => unify(tail, s.compose(new SimpleSubst(x, y)))
        case Some(t) => unify((y, t) :: tail, s)
      }
      case Some(conflict) => throw TypeError("Could not unify: " + y + " with " + conflict)
    }

    c match {
      case (TypeVar(a), TypeVar(b)) :: tail if (a == b) => unify(tail, s)
      case (TypeNat, TypeNat) :: tail => unify(tail, s)
      case (TypeBool, TypeBool) :: tail => unify(tail, s)
      case (TypeFun(a1, b1), TypeFun(a2, b2)) :: tail => unify((a1, a2) :: (b1, b2) :: tail, s)
      case (x: TypeVar, y) :: tail if notIn(x, y) => varSub(x, y, tail)
      case (x, y: TypeVar) :: tail if notIn(y, x) => varSub(y, x, tail)
      case (t1, t2) :: tail => throw TypeError("Could not unify: " + t1 + " with " + t2)
      case Nil => s
    }
  }

  override def typeOf(t: Term): Type = try {
    val TypingResult(tp, c) = collect(Nil: Env, t)
    val s = unify(c)
    s(tp)
  } catch {
    case TypeError(msg) =>
      Console.println("type error: " + msg)
      null
  }

}
